#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__ = "Lei Chen <hide1713@gmail.com>"
import operator
import StringIO
import urllib
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from numpy import array,log
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D
from matplotlib.figure import Figure
    
def save_and_display(fig):
    "Save figure to a image string"
    image_file = StringIO.StringIO()
    fig.savefig(image_file,format="png")
    return 'data:image/png,'+urllib.quote(image_file.getvalue())
    

def draw_segment_recruitment(model_list):
    """Draw segment recruitment
    model_list: a list of (cog_object, file_name)
    """
    fig = Figure()
    canvas = FigureCanvas(fig)
    number_of_row = len(model_list)
    extra_space = 30
    extra_space_y = 5
    for fig_index, model in enumerate(model_list):
        x_data = []
        e_vals = []
        range_list = model.get_hits_range()
        range_size = len(range_list)
        ax = fig.add_subplot(number_of_row, 1, fig_index + 1)
        ax.set_xlabel("Cog Length")
        ax.set_ylabel("Number of Reads")
        ax.set_xlim( -extra_space, model.length + extra_space)
        ax.set_ylim(0, range_size+1)
        ax.set_title(model.name + ":" + model.file_name)
        sorted(range_list, key = operator.itemgetter(3))
        for y, r in enumerate(reversed(range_list)):
            #r[0] is the read name of genome 
            #r[1] is the start position of genome read
            #r[2] is the end position of genome read
            #r[3] is the e value
            assert(r[1] >= 0 and r[2] <= model.length)
            x_data.append([(r[1], y + 1), (r[2], y + 1)])
            e_vals.append(r[3])
        base_line = Line2D([0, model.length],[y + 1, y + 1],\
                           linestyle = 'solid',linewidth = 5.0)
        arr = array(x_data)
        e_arr = array([log(e_val) for e_val in e_vals])
        line_segments =  LineCollection(arr,linestyles = 'solid')
        ax.add_collection(line_segments)
        ax.add_line(base_line)
        line_segments.set_array(e_arr)
        cb = fig.colorbar(line_segments, orientation = 'horizontal')
        cb.set_label("E Value (log)")
    return save_and_display(fig)
    
def draw_freq(cog_list):
    """Draw the frequency graph
    cog_list: is a list of Cog
    """
    from matplotlib.font_manager import FontProperties
    font = FontProperties(size = 'x-small')
    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(1,1,1)
    y_lim = []
    x_lim = []
    # Draw cog data
    hitp_data = cog_list[0].hitp
    #if hitp_data:
    #    ax.subplot(211)
    for cog in cog_list:
        freq = cog.get_hits_frequency()
        y_lim.append(max(freq))
        x_lim.append(cog.length)
        ax.plot(freq, label = cog.name + " " + cog.file_name)
    ax.set_ylim(0, max(y_lim))
    ax.set_xlim(0, max(x_lim))
    ax.set_xlabel("Length")
    ax.set_ylabel("Read Counts")
    ax.legend(prop = font)
    ax.set_title("Coverage frequency")
    ax.grid(True)
    # Draw hitp,  we assume that cogs are in the same family.
    if hitp_data:
        ax.twinx()
        ax.set_xlim(0, max(hitp_data))
        ax.plot(hitp_data, label = "Hitp", color = 'r')
    ax.legend(prop = font)
    return save_and_display(fig)   
 

def draw_function_category(models):
    """Draw function category"""
    from numpy import arange
    from random import choice
    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(1,1,1)
    category = set()
    fc_category = []
    #Read category from models
    for model in models:
        fc = model.get_function_category()
        fc_category.append(fc)
        for function_name in fc.func_cat.keys():
            category.add(function_name)
    category_list = [i for i in category]
    ax.set_xticks(arange(len(category_list)),category_list)
    y_max =  0
    ax.set_xlabel("function category")
    ax.set_ylabel("relative frequency")
    ax.set_title("Relative abundance(read-count based)")
    #Draw all the category
    for i, fc in enumerate(fc_category):
        values = []
        for cat_name in category_list:
            values.append(fc.func_cat.get(cat_name, 0))
        total_reads =  sum(values)
        #normalize the value
        values = [float(v)/float(total_reads) for v in values]
        ax.scatter(arange(len(category_list)), values,\
                    marker = "^s<>dph8+xo"[i%10], s = 25, \
                    c = "rgbcmykw"[i%8], label = fc.file_name)
        values.append(y_max)
    ax.legend()
    return save_and_display(fig)

def draw_similarity(models):
    import Pycluster
    import os
    from newick.tree import Tree, Leaf
    from drawtree.measure import Measure
    from drawtree.canvas  import PostScriptCanvas
    from drawtree.draw    import LinearDraw, BoxDraw
    data_arr = []
    data_name = []
    for a_model in models:
        data_name.append(a_model.base_name)
        data_arr.append(a_model.get_total_cog_names_and_hits())
    #Get the tree structure from Pycluster
    tree = Pycluster.treecluster(data = data_arr, dist = 's')
    #Translate tree to newick format
    g = Tree()
    sub_tree_map = {}
    for i, t in enumerate(tree):
        new_node =  Tree()
        #print t
        #add left node
        if t.left >= 0 :
            new_node.add_edge((Leaf(data_name[t.left]), None, 0.03))
        else:
            cluster_name = t.left
            new_node.add_edge((sub_tree_map[cluster_name], None, 0.03))
        #add right node
        if t.right >= 0:
            new_node.add_edge((Leaf(data_name[t.right]), None, 0.03))
        else:
            cluster_name = t.right
            new_node.add_edge((sub_tree_map[cluster_name], None, 0.03))
        sub_tree_map[ -i - 1] = new_node
    #Draw the tree into eps
    #here is a hack. The file is saved though src/drawtree/camvas.print_func
    if os.path.exists("temp.eps"):
        os.remove("temp.eps")
    g = sub_tree_map[min(sub_tree_map.keys())]
    mes = Measure()
    g.dfs_traverse(mes)
    canvas = PostScriptCanvas(g.length, g.leaves, mes.longest_label)
    draw = BoxDraw(canvas)
    canvas.begin()
    g.dfs_traverse(draw)
    canvas.end()
    os.system("epstopdf temp.eps")
    os.system("evince temp.pdf")







